# -*- coding: utf-8 -*-
'''
Created on Aug 6, 2019

@author: yl
'''

from FFT_Interpolation import *
from mpl_toolkits.mplot3d import Axes3D
from scipy import signal
from scipy.optimize import curve_fit
from scipy.signal import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation


j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]

pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size

w_0 = 1e-3
Z_R = np.pi * w_0**2 / Lamda

def fit_func(x, a, b):
    return a*x+b

def w_z(w_0, Z_R, Z):
    return w_0*np.sqrt(1 + (Z/Z_R)**2)

def R_z(Z_R, Z):
    return Z*(1+(Z_R/Z)**2)

def Phi_Gouy(Z_R, Z):
    return np.arctan(Z/Z_R)    

def f(X, Y, V_r_x, V_r_y, V_m_x, V_m_y, M=0.0165, N=0.033, L=0.12, D=0, w_0=w_0, Z_R=Z_R):
    
    k = 2 * np.pi / Lamda
    I_0 = 512
    
    Z_p_r = M + N + L * (2/(1-V_r_x**2-V_r_y**2) - 1)
    Z_p_m = M + N + (L+D) * (2/(1-V_m_x**2-V_m_y**2) - 1)
    
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
    
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    
    r_mea = np.sqrt( X**2 + Y**2 + (L+D)**2 - ((2*X*V_m_x+2*Y*V_m_y+(L+D)*(1-V_m_x**2-V_m_y**2))/(1+V_m_x**2+V_m_y**2))**2 )
    r_ref = np.sqrt( X**2 + Y**2 + (L)**2 - ((2*X*V_r_x+2*Y*V_r_y+L*(1-V_r_x**2-V_r_y**2))/(1+V_r_x**2+V_r_y**2))**2 )
    
    R_ref = R_z(Z_R, Z_p_r+d_ref)
    R_mea = R_z(Z_R, Z_p_m+d_mea)
    
    Phi_gouy_ref = Phi_Gouy(Z_R, Z_p_r+d_ref)
    Phi_gouy_mea = Phi_Gouy(Z_R, Z_p_m+d_mea)
    
    diff_phi = k * (diff_Z_p + (d_mea-d_ref) + r_mea**2/2/R_mea - r_ref**2/2/R_ref) + Phi_gouy_mea - Phi_gouy_ref
    #     A_beat = 0.5 * I_0 * w_0**2 / w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea) * np.exp((-X**2-Y**2)/w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea))
    I_beat = 1/w_z(w_0, Z_R, Z_p_r+d_ref)**2*np.exp(-2*r_ref**2/w_z(w_0, Z_R, Z_p_r+d_ref)**2) + 1/w_z(w_0, Z_R, Z_p_m+d_mea)**2*np.exp(-2*r_mea**2/w_z(w_0, Z_R, Z_p_m+d_mea)**2) + 2/w_z(w_0, Z_R, Z_p_r+d_ref)/w_z(w_0, Z_R, Z_p_m+d_mea)*np.exp(-r_ref**2/w_z(w_0, Z_R, Z_p_r+d_ref)**2-r_mea**2/w_z(w_0, Z_R, Z_p_m+d_mea)**2)*np.cos(diff_phi)
    I_beat = 0.5*I_0*w_0**2*I_beat
#     I_beat = I_beat.astype(np.int)
    return I_beat

def center_S(V_r_x=0, V_r_y=0, V_m_x=0, V_m_y=0, L=0.12,D=0):
    center_ref = [2*L*V_r_x/(1-V_r_x**2-V_r_y**2),2*L*V_r_y/(1-V_r_x**2-V_r_y**2)]
    center_mea = [2*(L+D)*V_m_x/(1-V_m_x**2-V_m_y**2),2*(L+D)*V_m_y/(1-V_m_x**2-V_m_y**2)]
#     print(center_ref, center_mea)
    return center_ref, center_mea

screen_diameter = 5e-3
dx = np.linspace(-screen_diameter/2, screen_diameter/2, num=pix_num)
dy = dx
X, Y = np.meshgrid(dx, dy)

V_r = 2e-3
V_r_x, V_r_y = V_r, V_r
V_m_x, V_m_y, D = 0, 0, 0
M, N, L = 0.0165, 0.033, 0.12
img = f(X, Y, V_r_x=V_r_x, V_r_y=V_r_y, V_m_x=V_m_x, V_m_y=V_m_y, M=M, N=N, L=L, D=D, w_0=w_0, Z_R=Z_R)

fig = plt.figure()
im = plt.imshow(img,cmap='gray',origin='lower', animated=True)
centers = center_S(V_r_x=V_r_x, V_r_y=V_r_y, V_m_x=V_m_x, V_m_y=V_m_y, L=L,D=D)
plt.colorbar(im, fraction=0.046, pad=0.04)
print(centers)

line, = plt.plot((img[640]/4+50), color='white')
line_ver, = plt.plot(img[:,640]/4+50, np.arange(1280), color='white')
center_ref = plt.scatter(centers[0][0]/pix_size+640,centers[0][1]/pix_size+640,color='r',marker='x')
center_mea = plt.scatter(centers[1][0]/pix_size+640,centers[1][1]/pix_size+640,color='b',marker='o')
num = 0 
def updatefig(*args):
    global D, num, V_m_x, V_m_y, V_r_x, V_r_y, L
    num = num % 10
#     D += -Lamda / 32
    V_m_x += 0.00005 * num
    #V_m_x = -0.013 - 0.00005 * num
    img = f(X, Y, V_r_x=V_r_x, V_r_y=V_r_y, V_m_x=V_m_x, V_m_y=V_m_y, M=M, N=N, L=L, D=D, w_0=w_0, Z_R=Z_R)
    centers = center_S(V_r_x=V_r_x, V_r_y=V_r_y, V_m_x=V_m_x, V_m_y=V_m_y, L=L,D=D)
    im.set_array(img)
    line.set_ydata(img[640]/4+50)
    line_ver.set_xdata(img[:,640]/4+50)
    center_ref.set_offsets([centers[0][0]/pix_size+640,centers[0][1]/pix_size+640])
    center_mea.set_offsets([centers[1][0]/pix_size+640,centers[1][1]/pix_size+640])
    num = num+1
    return im, line, line_ver,center_ref,center_mea

ani = animation.FuncAnimation(fig, updatefig, interval=50, blit=True)
plt.show()
# ani.save('test.gif', writer=animation.PillowWriter(fps=25))
# ani.save('test.gif')